package chat

import (
	"bigmodel/service/chat/spark"
	"context"
	"errors"

	"github.com/elliotchance/pie/v2"
	"github.com/zeromicro/go-zero/core/stores/cache"
	"github.com/zeromicro/go-zero/core/syncx"
)

type ResourcePool struct {
	syncx.Pool
	lim syncx.Limit
}

const (
	ResourceSparkPool = "spark"
	ResourceAliYunPool = "aliyun"
)

var Resource map[string]*ResourcePool

type ChatResourceFactoryParams struct {
	Sparks []spark.Client
	AKeys []string
	Cache  cache.Cache
}

func InitPool(params *ChatResourceFactoryParams) error {
	var size int
	r := make(map[string]*ResourcePool)
	if len(params.Sparks) > 0 {
		size = len(params.Sparks)
		r[ResourceSparkPool] = &ResourcePool{
			Pool: *syncx.NewPool(size, func() any {
				item := pie.Pop(&params.Sparks)
				return MustNewSpark(item.GetAppId(), item.GetApiKey(), item.GetApiSecret())
			}, func(a any) {}),
			lim: syncx.NewLimit(size),
		}
	}
	if len(params.AKeys) > 0 {
		size = len(params.AKeys)
		r[ResourceAliYunPool] = &ResourcePool{
			Pool: *syncx.NewPool(size,func () any {
				item := pie.Pop(&params.AKeys)
				return MustNewAliYun(*item)
			},func(a any){}),
			lim: syncx.NewLimit(size),
		}
	}
	Resource = r
	return nil

}

func GetResource(ctx context.Context, name string) (IChat, error) {
	if r, ok := Resource[name]; ok {
		if r.lim.TryBorrow() {
			return r.Get().(IChat), nil
		} else {
			return nil, errors.New("too many requests")
		}
	}
	return nil, errors.New("no such resource")
}

func ReleaseResource(ctx context.Context, name string, data any) error {
	if r, ok := Resource[name]; ok {
		r.Put(data)
		return r.lim.Return()
	}
	return nil
}
